mod ast;
mod compiler;
mod lexer;
mod parser;
mod passes;
mod rust_side;
mod token;

use std::error::Error;

use clap::{Parser as ClapParser, ValueEnum};
use compiler::Compiler;
use lexer::read_content;
use melior::{
    Context,
    dialect::DialectRegistry,
    ir::{Module, operation::OperationLike},
    pass::{
        PassManager,
        transform::{create_canonicalizer, create_cse, create_inliner},
    },
};
use parser::Parser;
use rust_side::{
    mlir_dump_llvm_ir, mlir_func_register_all_extensions, mlir_llvm_register_inliner_interface,
    mlir_load_toy_dialect, mlir_run_jit,
};

#[derive(ClapParser, Debug)]
#[command(version, long_about = None, about="toy compiler")]
struct Args {
    #[arg(
        value_enum,
        long,
        default_value_t = Action::Mlir,
        help = "Select the kind of output desired",
        value_name = "actin"
    )]
    emit: Action,

    #[arg(
        value_enum,
        short,
        default_value_t = InputType::Toy,
        help = "Decided the kind of output desired",
        value_name = "input type",
    )]
    x: InputType,

    #[arg(long, help = "Enable optimizations")]
    opt: bool,

    #[arg(index = 1, value_name = "input toy file", default_value_t = {"-".to_string()})]
    filename: String,
}

#[derive(ValueEnum, Clone, Debug, PartialEq, PartialOrd)]
enum Action {
    #[clap(help = "output the AST dump")]
    Ast,
    #[clap(help = "output the MLIR dump")]
    Mlir,
    #[clap(
        help = "output the MLIR dump after affine lowering",
        name = "mlir-affine"
    )]
    MlirAffine,
    #[clap(help = "output the MLIR dump after llvm lowering", name = "mlir-llvm")]
    MlirLLVM,
    #[clap(help = "output the LLVM IR dump", name = "llvm")]
    LLVMIR,
    #[clap(
        help = "JIT the code and run it by invoking the main function",
        name = "jit"
    )]
    RunJit,
}

#[derive(ValueEnum, Clone, Debug, PartialEq)]
enum InputType {
    #[clap(help = "load the input file as a Toy source")]
    Toy,
    #[clap(help = "load the input file as an MLIR file")]
    Mlir,
}

fn main() -> Result<(), Box<dyn Error>> {
    let cli = Args::parse();

    match cli.emit {
        Action::Ast => {
            dump_ast(&cli)?;
        }
        _ => {
            let registry = DialectRegistry::new();
            mlir_func_register_all_extensions(&registry);
            mlir_llvm_register_inliner_interface(&registry);
            let context = Context::new();
            context.append_dialect_registry(&registry);
            context.load_all_available_dialects();
            mlir_load_toy_dialect(&context);

            let module = load_and_process_mlir(&cli, &context)?;

            let is_outputing_mlir = cli.emit <= Action::MlirLLVM;
            if is_outputing_mlir {
                module.as_operation().dump();
                return Ok(());
            }

            if cli.emit == Action::LLVMIR {
                if !mlir_dump_llvm_ir(&context, &module, cli.opt) {
                    return Err("dump llvm ir error".into());
                }
                return Ok(());
            }

            if cli.emit == Action::RunJit {
                if !mlir_run_jit(&module, cli.opt) {
                    return Err("run jit error".into());
                }
                return Ok(());
            }
        }
    }
    Ok(())
}

fn load_and_process_mlir<'c>(
    cli: &Args,
    context: &'c Context,
) -> Result<Module<'c>, Box<dyn Error>> {
    let mut module = load_mlir(cli, &context)?;

    let pm = PassManager::new(&context);

    let is_lowering_to_affine = cli.emit >= Action::MlirAffine;
    let is_lowering_to_llvm = cli.emit >= Action::MlirLLVM;
    if cli.opt || is_lowering_to_affine {
        pm.add_pass(create_inliner());
        let op_pm = pm.nested_under("toy.func");
        op_pm.add_pass(passes::create_shape_inference());
        op_pm.add_pass(create_canonicalizer());
        op_pm.add_pass(create_cse());
    }

    if is_lowering_to_affine {
        pm.add_pass(passes::toy_create_lower_to_affine());

        let op_pm = pm.nested_under("func.func");
        op_pm.add_pass(create_canonicalizer());
        op_pm.add_pass(create_cse());

        if cli.opt {
            op_pm.add_pass(passes::affine_create_loop_fusion());
            op_pm.add_pass(passes::affine_create_affine_scalar_replacement());
        }
    }

    if is_lowering_to_llvm {
        pm.add_pass(passes::toy_create_lower_to_llvm());
        pm.add_pass(passes::llvm_craete_di_scope_for_llvm_func_op());
    }

    pm.run(&mut module)?;

    Ok(module)
}

fn dump_ast(cli: &Args) -> Result<(), Box<dyn Error>> {
    match cli.x {
        InputType::Toy => Parser::new(&cli.filename)
            .and_then(|mut p| p.parse())
            .map(|m| {
                if m.functions.is_empty() {
                    eprintln!("Parse error: empty module")
                } else {
                    println!("{}", m)
                }
            })?,
        InputType::Mlir => {
            return Err("Can't dump a Toy AST when the input is MLIR".into());
        }
    }
    Ok(())
}

fn load_mlir<'c>(cli: &Args, context: &'c Context) -> Result<Module<'c>, Box<dyn Error>> {
    Ok(
        if cli.x != InputType::Mlir && !cli.filename.ends_with(".mlir") {
            Compiler::new().compile(context, &cli.filename)?
        } else {
            let data = read_content(&cli.filename)?;
            Module::parse(&context, &data).ok_or::<Box<dyn Error>>("can't load file".into())?
        },
    )
}
